/*
 *  PktTest.cpp
 *  TubeTruck
 *
 *  Created by Aeturnum on 6/8/08.
 *  Copyright 2008 __MyCompanyName__. All rights reserved.
 *
 */

#include "Packet.h"
#include "OHAI.h"
#include "NOWAI.h"
#include "IIS.h"
#include "CANHAS.h"
#include "LAWL.h"
#include "DOWANT.h"
#include "HASFLAVR.h"
#include "CHEEZBURGER.h"
#include <vector>
#include "TubeTruckTypes.h"
#include "util.h"
#include "RSA.h"
#include "BlowFish.h"
#include "Signiture.h"
#include <stdlib.h>
#include <time.h>

using namespace TTP;
using namespace std;

char t[] = "True";
char f[] = "False";

char* getTruthString(bool value) {
   if(value)
      return t;
   return f;
}

void corruptPacket(Packet& p, vector<byte>& v) {
   try {  
      v[rand() % v.size()] = v[rand() % v.size()]+rand()%500;
      vector<byte>::const_iterator itr = v.begin(), end = v.end();
      p.fromVec(itr, end);
      printf("\n\tException Expected because of corruption, but NOT THROWN!");
   } catch(msg_exception e) {
      printf("\n\tException Thrown as Expected on Packet Corruption");
   }
}

void packetLevelTests(Packet& a, Packet& b, Packet& blank,  vector<byte>& intermediary, unsigned short packetType, uint32_t idHash, bool hasSourceId) {
   // basic tests
   printf("\n\tVersion Number: in: %d, out: %d, expected: %d", a.getVersionNumber(), b.getVersionNumber(), TubeTruckVersion);
   printf("\n\tType: in: %d, out: %d, expected: %d", a.getType(), b.getType(), packetType);
   if(hasSourceId) {
      printf("\n\tSource idHash: in: %d, out: %d, expected: %d", a.getSourceIdHash(), b.getSourceIdHash(), idHash);
   
      try {
         a.setDestIdHash(idHash);
         printf("\n\tException NOT Thrown on set Source Hash - Expected to Throw!");
      } catch(msg_exception e) {
         try {
            a.getDestIdHash();
            printf("\n\tException NOT Thrown on get Source Hash - Expected to Throw!");
         } catch(msg_exception e) {
            printf("\n\tExceptions Thrown as Expected on set/get Source Hash");
         }
      }
   } else {
      printf("\n\tDest idHash: in: %d, out: %d, expected: %d", a.getDestIdHash(), b.getDestIdHash(), idHash);
      
      try {
         a.setSourceIdHash(idHash);
         printf("\n\tException NOT Thrown on set Source Hash - Expected to Throw!");
      } catch(msg_exception e) {
         try {
            a.getSourceIdHash();
            printf("\n\tException NOT Thrown on get Source Hash - Expected to Throw!");
         } catch(msg_exception e) {
            printf("\n\tExceptions Thrown as Expected on set/get Source Hash");
         }
      }
      
   }

}

void testOHAI() {
   printf("\nStarting OHAI Test:");
   uint32_t idHash = 1;
   OHAI test;
   OHAI test2;
   OHAI blank;
   vector<byte> temp;
   
   test.setSourceIdHash(idHash);
   test.toVec(temp);
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_OHAI, idHash, true);
   // corruption detection
   corruptPacket(blank, temp);
   
   printf("\nEnding OHAI test\n");
}

void testNOWAI() {
   printf("\nStarting NOWAI Test:");
   uint32_t idHash = 1;
   NOWAI test;
   NOWAI test2;
   NOWAI blank;
   vector<byte> temp;
   
   test.setDestIdHash(idHash);
   test.toVec(temp);
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_NOWAI, idHash, false);
   corruptPacket(blank, temp);
   
   printf("\nEnding NOWAI test\n");
}

void testIIS() {
   printf("\nStarting IIS Test:");
   uint32_t idHash = 1;
   uint32_t DestHash = 2;
   IIS test;
   IIS test2;
   IIS blank;
   vector<byte> temp;
   
   vector<byte> pubKey;
   vector< vector<byte> > addrs;
   vector< vector<byte> > pubKeys;
   vector<byte> appData;
   
   
   pubKey.insert(pubKey.end(), 128, rand()%255);
   for(int i = 0; i< 3; i++) {
      vector<byte> a;
      vector<byte> k;
      a.insert(a.end(), 16, rand()%255);
      k.insert(k.end(), 128, rand()%255);
      addrs.insert(addrs.end(), a);
      pubKeys.insert(pubKeys.end(), k);
   }
   appData.insert(appData.end(), 300, rand()%255);
   appData.insert(appData.end(), 300, rand()%255);
   appData.insert(appData.end(), 300, rand()%255);
   appData.insert(appData.end(), 300, rand()%255);


   
   test.setSourceIdHash(idHash);
   test.setDestIdHash(DestHash);
   test.setPubKey(pubKey);
   test.addMulticastGroup(addrs[0], pubKeys[0]);
   test.addMulticastGroup(addrs[1], pubKeys[1]);
   test.addMulticastGroup(addrs[2], pubKeys[2]);
   test.addAppData(appData);
   test.toVec(temp);

   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr,end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_IIS, idHash, true);
   corruptPacket(blank, temp);
   printf("\n\n\tShould all be true:");
   printf("\n\tOld and New Pubkey the same: %s", getTruthString(pubKey == test2.getPubKey()));
   printf("\n\tOld and New Dest Hash same: %s", getTruthString(DestHash == test2.getDestIdHash()));
   printf("\n\tOld and New Group[0] Key and Addr the same: %s", getTruthString(addrs[0] == test2.getMcastAddr(0) && pubKeys[0] == test2.getGroupKey(0)));
   printf("\n\tOld and New Group[1] Key and Addr the same: %s", getTruthString(addrs[1] == test2.getMcastAddr(1) && pubKeys[1] == test2.getGroupKey(1)));
   printf("\n\tOld and New Group[2] Key and Addr the same: %s", getTruthString(addrs[2] == test2.getMcastAddr(2) && pubKeys[2] == test2.getGroupKey(2)));
   printf("\n\tOld and New appData the same: %s", getTruthString(appData == test2.getAppData()));
   
   printf("\nEnding IIS test\n");
}

void testCANHAS() {
   printf("\nStarting CANHAS Test:");
   uint32_t idHash = 1;
   CANHAS test;
   CANHAS test2;
   CANHAS blank;
   vector<byte> temp;
   vector<byte> pubKey;
   vector<byte> groupPubKey;
   
   pubKey.insert(pubKey.end(), 128, rand()%255);
   groupPubKey.insert(groupPubKey.end(), 128, rand()%255);
   
   test.setSourceIdHash(idHash);
   test.setMyPublicKey(pubKey);
   test.setGroupPublicKey(groupPubKey);
   test.toVec(temp);
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_CANHAS, idHash, true);
   corruptPacket(blank, temp);
   printf("\n\n\tShould all be true:");
   printf("\n\tOld and New Pubkey the same: %s", getTruthString(pubKey == test2.getMyPublicKey()));
   printf("\n\tOld and New Group Pubkey the same: %s", getTruthString(groupPubKey == test2.getGroupPublicKey()));

   
   
   printf("\nEnding CANHAS test\n");
}

void testLAWL() {
   printf("\nStarting LAWL Test:");
   uint32_t idHash = 1;
   Crypt::RSA myRsa(1024);
   Crypt::RSA theirRSA;
   LAWL test;
   LAWL test2;
   LAWL blank;
   vector<byte> temp;
   vector<byte> pubKey;
   vector<byte> RSAKey = myRsa.getPublicKey();
   
   theirRSA.setPublicKey(RSAKey);
   
   pubKey.insert(pubKey.end(), 128, rand()%255);
   
   test.setDestIdHash(idHash);
   test.setGroupPublicKey(pubKey);
   myRsa.sign(test);
   test.toVec(temp);
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_LAWL, idHash, false);
   printf("\n\n\tShould all be true:");
   printf("\n\tSigniture should verify: %s", getTruthString(theirRSA.verify(test2)));
   printf("\n\tOld and New Group Pubkey the same: %s", getTruthString(pubKey == test2.getGroupPublicKey()));
   
   printf("\nEnding LAWL test\n");
}

void testHASFLAVR() {
   printf("\nStarting HASFLAVR Test:");
   uint32_t idHash = 1;
   Crypt::RSA myRsa(1024);
   Crypt::RSA theirRSA;
   
   int keyArr[16] = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};
   int ivArr[8] = {1,2,3,4,5,6,7,8};
   HASFLAVR test;
   HASFLAVR test2;
   HASFLAVR blank;
   vector<byte> key(keyArr, keyArr+16), iv(ivArr, ivArr+8);
   Crypt::BlowFish myBlowfish(key, iv);
   vector<byte> temp;
   vector<byte> RSAKey = myRsa.getPublicKey();
   
//   key = myBlowfish.getKey();
//   iv = myBlowfish.getIv();
   
   theirRSA.setPublicKey(RSAKey);
   
   test.setSourceIdHash(idHash);
   test.setFlag(HASFLAVR::HASFLAVR_no);
   myBlowfish.encrypt(test);
   iv = myBlowfish.getIv();
   test.setBlowfishIv(iv);
   myRsa.sign(test);
   test.toVec(temp);
   
//   cout << "\nTemp(" << temp.size() << "):\n" << temp;
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   vector<byte> iv2 = test2.getBlowfishIv();
   
   Crypt::BlowFish theirBlowfish(key, iv2);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_HASFLAVR, idHash, true);
   printf("\n\n\tShould all be true:");
   printf("\n\tSigniture should verify: %s", getTruthString(theirRSA.verify(test2)));
   theirBlowfish.decrypt(test2);
   printf("\n\tFlags should be the same: %s", getTruthString(test.getFlags() == test2.getFlags()));
   
   printf("\nEnding HASFLAVR test\n");
}

void testDOWANT() {
   printf("\nStarting DOWANT Test:");
   uint32_t idHash = 1;
   Crypt::RSA myGroupRsa(1024);
   Crypt::RSA myTheirRSA;
   Crypt::RSA theirRSA(1024);
   Crypt::RSA theirGroupRSA;
   DOWANT test;
   DOWANT test2;
   DOWANT blank;
   vector<byte> temp;
   vector<byte> priKey;
   vector<byte> groupSymKey;
   
   vector<byte> theirPubKey = theirRSA.getPublicKey();
   vector<byte> groupPubKey = myGroupRsa.getPublicKey();
   myTheirRSA.setPublicKey(theirPubKey);
   theirGroupRSA.setPublicKey(groupPubKey);
   
   priKey = myGroupRsa.getPrivateKey();
   
   groupSymKey.insert(groupSymKey.end(), 16, rand()%255);
   
   test.setDestIdHash(idHash);
   test.setSymmetricKey(groupSymKey);
   test.setGroupPrivateKey(priKey);
   myTheirRSA.encrypt(test);
   myGroupRsa.sign(test);
   test.toVec(temp);
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr, end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_DOWANT, idHash, false);
   printf("\n\n\tShould all be true:");
   printf("\n\tSigniture should verify: %s", getTruthString(theirGroupRSA.verify(test2)));
   theirRSA.decrypt(test2);
   printf("\n\tOld and New Group Sym Key the same: %s", getTruthString(groupSymKey == test2.getSymmetricKey()));
   printf("\n\tOld and New Group Private Key the same: %s", getTruthString(priKey == test2.getGroupPrivateKey()));
   
   printf("\nEnding DOWANT test\n");
}

void testCHEEZBURGER() {
   printf("\nStarting CHEEZBURGER Test:");
   uint32_t idHash = 1;
   CHEEZBURGER test;
   CHEEZBURGER test2;
   CHEEZBURGER blank;
   Crypt::RSA myRsa(1024);
   Crypt::RSA theirRSA;
   Crypt::BlowFish myBlowfish;
   vector<byte> temp;
   vector<byte> RSAkey;
   vector<uint32_t> idHashs;
   vector<uint32_t> theirAckNumbers;
   vector<uint32_t> myAckNumbers;
   vector< vector<byte> > data;
   
   
   RSAkey = myRsa.getPublicKey();
   theirRSA.setPublicKey(RSAkey);
   
   for(int i = 0; i< 3; i++) {
      vector<byte> d;
      
      idHashs.push_back(rand()%65535);
      theirAckNumbers.push_back(rand()%65535);
      myAckNumbers.push_back(rand()%65535);
      d.insert(d.end(), rand()%1000, rand()%255);
      data.insert(data.end(), d);
   }
   

   vector<byte> key = myBlowfish.getKey();
   test.setSourceIdHash(idHash);
   test.addHashAckPair(idHashs[0], theirAckNumbers[0]);
   test.addHashAckPair(idHashs[1], theirAckNumbers[1]);
   test.addHashAckPair(idHashs[2], theirAckNumbers[2]);
   test.setAckNumber(myAckNumbers[0]);
   test.addDataSegment(data[0]);
   test.addDataSegment(myAckNumbers[1], data[1]);
   test.setAckNumber(myAckNumbers[2]);
   test.addDataSegment(data[2]);
   myBlowfish.encrypt(test);
   vector<byte> iv = myBlowfish.getIv();
   test.setBlowfishIv(iv);
   myRsa.sign(test);
   test.toVec(temp);
   
//   cout << "temp(" << temp.size() << "):\n" <<temp ;
   
   vector<byte>::const_iterator itr = temp.begin(), end = temp.end();
   
   test2.fromVec(itr,end);
   
   packetLevelTests(test, test2, blank, temp, Packet::Packet_CHEEZBURGER, idHash, true);
   printf("\n\n\tShould all be true:");
   printf("\n\tSigniture should verify: %s", getTruthString(theirRSA.verify(test2)));
   fflush(stdout);
   vector<byte> transmittedIv = test2.getBlowfishIv();
   
   Crypt::BlowFish theirBlowfish(key, transmittedIv);
   theirBlowfish.decrypt(test2);
   printf("\n\tDecrypted!");
   fflush(stdout);
   
   vector<uint32_t> idHashs2 = test2.getHashList();
   vector<uint32_t> theirAckNumbers2 = test2.getAckNumbers();
   printf("\n\tOld and New Hash-ack pair[0]() the same: %s", getTruthString(idHashs[0] == idHashs2[0] && theirAckNumbers[0] == test2.getAckNumber(idHashs2[0]) && theirAckNumbers[0] == theirAckNumbers2[0]));
   fflush(stdout);
   printf("\n\tOld and New Hash-ack pair[1]() the same: %s", getTruthString(idHashs[1] == idHashs2[1] && theirAckNumbers[1] == test2.getAckNumber(idHashs2[1]) && theirAckNumbers[1] == theirAckNumbers2[1]));
   fflush(stdout);
   printf("\n\tOld and New Hash-ack pair[2]() the same: %s", getTruthString(idHashs[2] == idHashs2[2] && theirAckNumbers[2] == test2.getAckNumber(idHashs2[2]) && theirAckNumbers[2] == theirAckNumbers2[2]));
   fflush(stdout);
   vector<uint32_t> transmittedAckNums = test2.getDataAckNumbers();
   assert(transmittedAckNums.size() == 3);
   printf("\n\tOld and New data vectors[0] the same: %s", getTruthString(data[0] == test2.getDataVector(transmittedAckNums[0])));
   printf("\n\tOld and New data vectors[0] the same: %s", getTruthString(data[1] == test2.getDataVector(transmittedAckNums[1])));
   printf("\n\tOld and New data vectors[0] the same: %s", getTruthString(data[2] == test2.getDataVector(transmittedAckNums[2])));
   fflush(stdout);
   printf("\nEnding CHEEZBURGER test\n");
   
}


/*
 General Comments:
   Packet Classes:
      Hashing (throw exceptions on bad hash - no order what-so-ever to calls):
         OHAI
         NOWAI
         IIS
         CANHAS
      Signed: (DO NOT throw exceptions, RSA.verify(Packet&) returns FALSE on either corruption or tampering)
               MUST NOT MODIFY DATA AFTER SIGNING - ENCRYPT BEFORE SIGNING (if applicable)
         LAWL
         HASFLAVR
         DOWANT
         CHEEZBURGER
      Encrypted w/ RSA: (Encrypt BEFORE signing. Just use RSA.encrypt(Packet&) / RSA.decrypt(Packet&)) - VERIFY FIRST
         DOWANT
      Encrypted w/ BlowFish: (Encrypt BEFORE signing. Just use BlowFish.encrypt(Packet&) / BlowFish.decrypt(Packet&)) - VERIFY FIRST
                              BLOWFISH objects change their keys BEFORE EACH encryption, observe the following calling convention:
                              Packet p;
                              BlowFish.Encrypt(p);
                              vector<byte> iv = BlowFish.getIv();
                              p.setBlowfishIv(iv);
                              RSA.sign(p);
                              -- ANY other order will probably FUBAR the packet and maybe crash the program --
         HASFLAVR
         CHEEZBURGER
 */

int main(int argc, char* argv[]) {
   srand(clock());
   testOHAI();
   testNOWAI();
   testIIS();
   testCANHAS();
   testLAWL();
   testHASFLAVR();
   testDOWANT();
   testCHEEZBURGER();
}